/*
********************************************************************************
*
*	file : update
*
********************************************************************************
*/
#include "dataproc.h"
#include "stdio.h"
#include "mytype.h"
#include "swap.h"
#include "mbcom.h"
#include "str_ext.h"
#include "app_update.h"
#include "aes.h"
#include "md5.h"
#include "sha256.h"
#include "common.h"

/********************************************************************************
*                               cfg
********************************************************************************/
/*******************************************************************************
* @brief  文件组织格式
* 00-03 : 文件数据区位置
* 04-07 : 文件数据区长度
* 08-0B : AES 模式
* 20-2F : 文件数据区MD5校验值
* 40-5F : 文件数据区SHA校验值
* 无论是否开了AES，正文长度强制为16的倍数
* 正文长度至少1KB
*******************************************************************************/
#define CODE_PAGE_SIZE 1024UL

#define APP_FLASH_BASE iap_obj.app_mem_base
#define APP_FLASH_SIZE iap_obj.app_mem_size
#define APP_ID_MEM_OFFSET iap_obj.id_mem_offset
#define APP_ID_MEM_VALUE iap_obj.id_mem_value

/********************************************************************************
*                               fun
********************************************************************************/
void enc_iap_opt(int opt);
void enc_iap_info_out(void);
void enc_iap_flash_end(void);
void enc_iap_flash(uint32_t offset, uint16_t len);
void enc_iap_dat_dec(uint16_t len);
void enc_iap_dec_init(uint16_t aes_bits);
void enc_iap_head_check(void);
void enc_iap_hash_compare(void);

uint8_t is_enc_iap_flash_different(uint32_t offset, uint16_t len);

/********************************************************************************
*                               var
********************************************************************************/
mbedtls_aes_context *aes_obj;
mbedtls_sha256_context *sha_src;
mbedtls_md5_context *md5_src;

// 此数组为 key_code_sim 代码二进制值
// 作为固定键值，可以伪装为程序
const uint16_t iap_ase_key_src[] = {
	0xf04f, 0x5000, 0x6801, 0x07ca,
	0xd500, 0xbf00, 0x6841, 0x078a,
	0xd500, 0xbf00, 0x6881, 0x074a,
	0xd500, 0xbf00, 0x68c0, 0x0701,
	0xd500, 0xbf00, 0x4770};

const uint8_t iap_ase_key_ext[32] = {
	0x00};

/*******************************************************************************
* @brief  main : std proc entry.
* \param[in] none
* \retval: none
*******************************************************************************/
void enc_iap_update(uint8_t check0_or_program1)
{
	int cnt;

	ap_ext_start();
	iap_obj.out.st = 0;

	// 申请内存
	aes_obj = malloc(sizeof(mbedtls_aes_context) + 16);
	sha_src = malloc(sizeof(mbedtls_aes_context) + 16);
	md5_src = malloc(sizeof(mbedtls_aes_context) + 16);

	//升级文件检测
	cnt = 3;
	while (cnt--)
	{
		enc_iap_opt(0);

		if (!iap_obj.err.st)
		{
			break;
		}
	}

	//升级
	if (check0_or_program1)
	{
		iap_obj.out.bits.update_req = 1;

		if (!iap_obj.err.st && iap_obj.out.bits.app_diff)
		{
			cnt = 3;
			while (cnt--)
			{
				enc_iap_opt(1);

				if (!iap_obj.err.st)
				{
					break;
				}
			}
		}
	}

	// 结果输出
	enc_iap_info_out();

	// 释放内存
	free(aes_obj);
	free(sha_src);
	free(md5_src);

	ap_ext_end();
}

/*******************************************************************************
* @brief  enc_iap_opt
* \param[in] write_en: 0-juat check, 1-check and write
* \retval: none
*******************************************************************************/
void enc_iap_opt(int write_en)
{
	int index;
	uint32_t len_read_total;
	uint16_t len_dec;

	iap_obj.err.st = 0;

	//检测程序是否有效
	do
	{
		//打开文件
		if (!ap_file_open(iap_obj.src_file, FIL_OPEN_R))
		{
			iap_obj.err.bits.err_no_file = 1;
			break;
		}

		// 文件头检测
		enc_iap_head_check();
		if (iap_obj.err.st)
		{
			break;
		}
		iap_obj.enc_file_size = ap_file_size_get();
		iap_obj.app_bin_size = iap_obj.hash_iap_head.app_len;
		if (iap_obj.app_bin_size > APP_FLASH_SIZE)
		{
			iap_obj.err.bits.err_size = 1;
			break;
		}

		// 初始化解密器
		enc_iap_dec_init(iap_obj.hash_iap_head.aes_bit);

		//读取检测
		len_read_total = 0;
		ap_file_seek(iap_obj.hash_iap_head.app_site);

		for (; len_read_total < iap_obj.hash_iap_head.app_len; len_read_total += len_dec)
		{
			memset(iap_obj.buff, 0xff, CODE_PAGE_SIZE);

			// 单次解码长度计算
			if (len_read_total + CODE_PAGE_SIZE <= iap_obj.hash_iap_head.app_len)
				len_dec = CODE_PAGE_SIZE;
			else
				len_dec = iap_obj.hash_iap_head.app_len - len_read_total;

			//读取一页数据
			if (0 == ap_file_read(iap_obj.buff, len_dec))
			{
				iap_obj.err.bits.err_file_read = 1;
				break;
			}

			//密文解算
			enc_iap_dat_dec(len_dec);

			// 程序ID
			if ((len_read_total == 0) && (APP_ID_MEM_OFFSET))
			{
				uint32_t *p_id;
				p_id = (uint32_t *)&iap_obj.buff[APP_ID_MEM_OFFSET];

				if (*p_id != APP_ID_MEM_VALUE)
				{
					iap_obj.err.bits.id_err = 1;
				}
			}

			//flash编程
			if (write_en)
			{
				if (!iap_obj.err.st)
				{
					enc_iap_flash(len_read_total, len_dec);
				}
			}
			else
			{
				if (is_enc_iap_flash_different(len_read_total, len_dec))
				{
					iap_obj.out.bits.app_diff = 1;
				}
			}
		}

		if (iap_obj.err.st)
		{
			break;
		}

		// hash 验证
		enc_iap_hash_compare();

	} while (0);

	ap_file_close();

	if (iap_obj.err.st)
	{
		iap_obj.out.bits.app_diff = 1;
	}

	//升级完成操作
	if (iap_obj.err.st == 0)
	{
		if (write_en)
		{
			enc_iap_flash_end();
		}
	}
}

/*******************************************************************************
* @brief  enc_iap_dec_init : 解码 初始化
* \param[in] none
* \retval: none
*******************************************************************************/
void enc_iap_dec_init(uint16_t aes_bits)
{
	uint8_t aes_key[48];

	// hash
	mbedtls_sha256_init(sha_src);
	mbedtls_sha256_starts(sha_src, 0);

	mbedtls_md5_init(md5_src);
	mbedtls_md5_starts(md5_src);

	// 192 BITS AES KEY => 24 BYTES
	mbedtls_aes_init(aes_obj);
	copy_u8(iap_ase_key_src, aes_key, 32);
	swap_bit_map_u16(aes_key, aes_key);
	map_add_u8(aes_key, (uint8_t *)iap_ase_key_ext, aes_key, 32);

	switch (aes_bits)
	{
	case 128:
		mbedtls_aes_setkey_dec(aes_obj, aes_key, 128);
		break;
	case 192:
		mbedtls_aes_setkey_dec(aes_obj, aes_key, 192);
		break;
	case 256:
		mbedtls_aes_setkey_dec(aes_obj, aes_key, 256);
	default:
		break;
	}
}

/*******************************************************************************
* @brief  update_dat_dec : 解码
* \param[in] none
* \retval: none
*******************************************************************************/
void enc_iap_dat_dec(uint16_t len)
{
	int index;
	uint8_t dec_buff[18];

	if (iap_obj.hash_iap_head.aes_mode)
	{
		for (index = 0; index < len; index += 16)
		{
			mbedtls_aes_crypt_ecb(aes_obj, MBEDTLS_AES_DECRYPT, &iap_obj.buff[index], dec_buff);
			copy_u8(dec_buff, &iap_obj.buff[index], 16);
		}
	}

	mbedtls_sha256_update(sha_src, iap_obj.buff, len);
	mbedtls_md5_update(md5_src, iap_obj.buff, len);
}

/*******************************************************************************
* @brief  enc_iap_compare : hash 比对
* \param[in] none
* \retval: none
*******************************************************************************/
void enc_iap_hash_compare(void)
{
	uint8_t hash_buff[36];

	// sha 验证
	mbedtls_sha256_finish(sha_src, iap_obj.buff);
	ap_file_seek(0x40);
	ap_file_read(iap_obj.comm_buff, 32);
	if (comp_i32(iap_obj.buff, iap_obj.comm_buff, 32 / 4))
		iap_obj.err.bits.sha = 1;

	// md5 验证
	mbedtls_md5_finish(md5_src, iap_obj.buff);
	ap_file_seek(0x20);
	ap_file_read(iap_obj.comm_buff, 16);
	if (comp_i32(iap_obj.buff, iap_obj.comm_buff, 16 / 4))
		iap_obj.err.bits.md5 = 1;
}

/*****************************************************************************
* @brief   simp_fat_update_head_check.
* @param   none
* @return  none
*****************************************************************************/
void enc_iap_head_check(void)
{
	uint32_t len_back;
	uint32_t buff[6];
	
	// 读取信息文件
	ap_file_seek(0);
	len_back = ap_file_read(buff, 18);
	iap_obj.hash_iap_head.app_site = buff[0];
	iap_obj.hash_iap_head.app_len = buff[1];
	iap_obj.hash_iap_head.aes_mode = byte_lb(buff[4]);
	
	if (!len_back)
	{
		iap_obj.err.bits.err_file_read = 1;
		return;
	}

	// 信息验证
	if (iap_obj.hash_iap_head.app_site < 0x80)
	{
		dbg_puts("fil main site error !\n");
		iap_obj.err.bits.err_size = 1;
		return;
	}

	if (ap_file_size_get() < iap_obj.hash_iap_head.app_site + iap_obj.hash_iap_head.app_len)
	{
		dbg_puts("fil info site error !\n");
		iap_obj.err.bits.err_size = 1;
		return;
	}

	// AES mode
	switch (iap_obj.hash_iap_head.aes_mode)
	{
	case 0x33:
		iap_obj.hash_iap_head.aes_bit = 128;
		break;
	case 0x55:
		iap_obj.hash_iap_head.aes_bit = 192;
		break;
	case 0x77:
		iap_obj.hash_iap_head.aes_bit = 256;
		break;
	default:
		iap_obj.hash_iap_head.aes_bit = 0;
		iap_obj.err.bits.aes_mode = 0;
		dbg_puts("iap enc mode error !\n");
		break;
	}
}

/*******************************************************************************
* @brief  main : 检测一页数据是否完全相同
* \param[in] offset 偏移地址
* \retval: diff
*******************************************************************************/
uint8_t is_enc_iap_flash_different(uint32_t offset, uint16_t len)
{
	//读取
	ap_flash_read(APP_FLASH_BASE + offset, len, iap_obj.comm_buff);

	return comp_i32(iap_obj.buff, iap_obj.comm_buff, len >> 2);
}

/*******************************************************************************
* @brief  main : update_flash. 编程一页数据
* \param[in] none
* \retval: none
*******************************************************************************/
void enc_iap_flash(uint32_t offset, uint16_t len)
{
	// 开始编程时擦除
	if (offset == 0)
	{
		copy_u32(iap_obj.buff, iap_obj.head_buff, CODE_PAGE_SIZE / 4);
		ap_flash_erase(APP_FLASH_BASE, APP_FLASH_BASE + iap_obj.app_bin_size - 1);
	}

	//首页不直接编程
	if (offset != 0)
	{
		ap_flash_write(APP_FLASH_BASE + offset, len, iap_obj.buff);

		if (is_enc_iap_flash_different(offset, len))
		{
			iap_obj.err.bits.verify = 1;
		}
	}
}

/*******************************************************************************
* @brief  main : update_flash_end.
* \param[in] none
* \retval: none
*******************************************************************************/
void enc_iap_flash_end(void)
{
	// 编程第一页
	ap_flash_write(APP_FLASH_BASE, CODE_PAGE_SIZE, iap_obj.head_buff);
	copy_u32(iap_obj.head_buff, iap_obj.buff, CODE_PAGE_SIZE / 4);

	if (is_enc_iap_flash_different(0, CODE_PAGE_SIZE))
	{
		iap_obj.err.bits.verify = 1;
	}
}

/*******************************************************************************
* @brief  main : update_info_out.
* \param[in] none
* \retval: none
*******************************************************************************/
void enc_iap_info_out(void)
{
	int index;

	// 无升级文件时不输出
	if (iap_obj.err.bits.err_no_file)
		return;

	if (ap_file_open(iap_obj.out_file, FIL_OPEN_WA))
	{
		ap_file_seek(ap_file_size_get());

		index = 0;

		index = sprintf((char *)&iap_obj.buff[0], "\r\n...........................................................\r\n");
		ap_file_write(iap_obj.buff, index);

		if (0 == iap_obj.err.st)
		{
			index = 0;

			if (iap_obj.out.bits.update_req)
			{
				// 检测并升级
				index += sprintf((char *)&iap_obj.buff[index], "Uptade success!\r\n");

				if (!iap_obj.out.bits.app_diff)
				{
					// BIN文件相同
					index += sprintf((char *)&iap_obj.buff[index], "The app file is the same as the input file! \r\n");
				}
				else
				{
					// 更新成功
					index += sprintf((char *)&iap_obj.buff[index], "Uptade file size %d Bytes!\r\n", iap_obj.enc_file_size);
					index += sprintf((char *)&iap_obj.buff[index], "Uptade bin  size %d Bytes!\r\n", iap_obj.app_bin_size);
				}
				
				iap_obj.out.bits.update_ok = 1;
			}
			else
			{
				// 仅检测
				index += sprintf((char *)&iap_obj.buff[index], "Check success!\r\n");
				if (iap_obj.out.bits.app_diff)
				{
					// 文件检查通过
					index += sprintf((char *)&iap_obj.buff[index], "Uptade file size %d Bytes!\r\n", iap_obj.enc_file_size);
					index += sprintf((char *)&iap_obj.buff[index], "Uptade bin  size %d Bytes!\r\n", iap_obj.app_bin_size);
				}
				else
				{
					// 文件相同
					index += sprintf((char *)&iap_obj.buff[index], "The app file is the same as the input file! \r\n");
				}
				iap_obj.out.bits.check_ok = 1;
			}
			
			ap_file_write(iap_obj.buff, index);
			dbg_puts((char *)iap_obj.buff);
		}
		else
		{
			if (iap_obj.out.bits.update_req)
			{
				iap_obj.out.bits.update_err = 1;
			}
			
			index = sprintf((char *)&iap_obj.buff[0], "UPDATE OR CHECK ERROR::\r\n");
			ap_file_write(iap_obj.buff, index);
			dbg_puts((char *)iap_obj.buff);
			
			if (iap_obj.err.bits.err_no_file)
			{
				index = sprintf((char *)&iap_obj.buff[0], "No file found! -> yt_app.bin \r\n");
				ap_file_write(iap_obj.buff, index);
			}

			if (iap_obj.err.bits.err_crc)
			{
				index = sprintf((char *)&iap_obj.buff[0], "The file crc is error! \r\n");
				ap_file_write(iap_obj.buff, index);
			}

			if (iap_obj.err.bits.err_sum)
			{
				index = sprintf((char *)&iap_obj.buff[0], "The app file is incomplete! \r\n");
				ap_file_write(iap_obj.buff, index);
			}

			if (iap_obj.err.bits.verify)
			{
				index = sprintf((char *)&iap_obj.buff[0], "The app flash verify fail! \r\n");
				ap_file_write(iap_obj.buff, index);
			}

			if (iap_obj.err.bits.err_size)
			{
				index = sprintf((char *)&iap_obj.buff[0], "The app file size is err! \r\n");
				ap_file_write(iap_obj.buff, index);
			}

			if (iap_obj.err.bits.err_file_read)
			{
				index = sprintf((char *)&iap_obj.buff[0], "The app file file read err! \r\n");
				ap_file_write(iap_obj.buff, index);
			}

			if (iap_obj.err.bits.aes_mode)
			{
				index = sprintf((char *)&iap_obj.buff[0], "File encode mode err! \r\n");
				ap_file_write(iap_obj.buff, index);
			}

			if (iap_obj.err.bits.md5)
			{
				index = sprintf((char *)&iap_obj.buff[0], "File MD5 err! \r\n");
				ap_file_write(iap_obj.buff, index);
			}

			if (iap_obj.err.bits.sha)
			{
				index = sprintf((char *)&iap_obj.buff[0], "File SHA err! \r\n");
				ap_file_write(iap_obj.buff, index);
			}

			if (!iap_obj.err.st && iap_obj.err.bits.id_err)
			{
				index = sprintf((char *)&iap_obj.buff[0], "The app program id is not match! \r\n");
				ap_file_write(iap_obj.buff, index);
			}
		}
	}

	index = sprintf((char *)&iap_obj.buff[0], "...........................................................\r\n");
	ap_file_write(iap_obj.buff, index);

	ap_file_close();
}
